// -----------------------------------------------------------------------
// <copyright file="MessageHeaderTests.cs" company="Asynkron AB">
//      Copyright (C) 2015-2024 Asynkron AB All rights reserved
// </copyright>
// -----------------------------------------------------------------------

using System;
using System.Threading.Tasks;
using FluentAssertions;
using Proto.Utils;
using Xunit;
using Xunit.Abstractions;

namespace Proto.Tests.Headers;

public class MessageHeaderTests
{
    private readonly ITestOutputHelper _testOutputHelper;

    public MessageHeaderTests(ITestOutputHelper testOutputHelper)
    {
        _testOutputHelper = testOutputHelper;
    }

    [Fact]
    public async Task HeadersArePropagatedBackInReply()
    {
        Sender PropagateHeaders(Sender next) =>
            (context, target, envelope) =>
                next(context, target, envelope.WithHeader(context.Headers));

        var headers = MessageHeader.Empty.With("foo", "bar");

        var system = new ActorSystem(ActorSystemConfig.Setup() with
        {
            ConfigureRootContext = context => context.WithHeaders(headers)
        });

        var props1 = Props.FromFunc(ctx =>
                {
                    switch (ctx.Message)
                    {
                        case SomeRequest:
                            ctx.Respond(new SomeResponse());

                            return Task.CompletedTask;
                        default:
                            return Task.CompletedTask;
                    }
                }
            )
            .WithSenderMiddleware(PropagateHeaders);

        var pid1 = system.Root.Spawn(props1);

        var tcs1 = new TaskCompletionSource<MessageHeader>();
        var tcs2 = new TaskCompletionSource<MessageHeader>();

        var props2 = Props.FromFunc(ctx =>
                {
                    switch (ctx.Message)
                    {
                        case StartMessage:
                            tcs1.SetResult(ctx.Headers);
                            ctx.Request(pid1, new SomeRequest());

                            break;
                        case SomeResponse:
                            tcs2.SetResult(ctx.Headers);

                            break;
                    }

                    return Task.CompletedTask;
                }
            )
            .WithSenderMiddleware(PropagateHeaders);

        var pid2 = system.Root.Spawn(props2);

        //ensure we set the headers up correctly
        Assert.Equal("bar", headers.GetOrDefault("foo"));

        //use the headers and send the request
        var root = system
            .Root
            .WithSenderMiddleware(PropagateHeaders);

        root.Send(pid2, new StartMessage());

        //actor1 should have headers
        var headers1 = await tcs1.Task.WithTimeout(TimeSpan.FromSeconds(5));
        Assert.Equal(headers, headers1);

        //actor2 should have headers
        var headers2 = await tcs2.Task.WithTimeout(TimeSpan.FromSeconds(5));
        Assert.Equal(headers, headers2);
    }

    [Fact]
    public async Task Actors_can_reply_with_headers()
    {
        await using var system = new ActorSystem();

        var echo = Props.FromFunc(ctx =>
            {
                if (ctx.Sender is not null && ctx.Message is not null)
                {
                    var messageHeader = MessageHeader.Empty.With("foo", "bar");
                    ctx.Respond(ctx.Message, messageHeader);
                }

                return Task.CompletedTask;
            }
        );

        var pid = system.Root.Spawn(echo);

        const int message = 1;
        var (msg, header) = await system.Root.RequestWithHeadersAsync<int>(pid, message);

        msg.Should().Be(message);
        header["foo"].Should().Be("bar");
    }

    [Fact]
    public async Task RequestAsync_honors_message_envelopes()
    {
        await using var system = new ActorSystem();

        var echo = Props.FromFunc(ctx =>
            {
                if (ctx.Sender is not null && ctx.Headers.Count == 1)
                {
                    ctx.Respond(ctx.Headers["foo"]);
                }

                return Task.CompletedTask;
            }
        );

        var pid = system.Root.Spawn(echo);

        var wrongPid = PID.FromAddress("some-incorrect-address", "some-id");

        var response = await system.Root.RequestAsync<string>(pid,
            new MessageEnvelope(1, wrongPid, MessageHeader.Empty.With("foo", "bar")),
            CancellationTokens.FromSeconds(1)
        );

        response.Should().Be("bar");
    }

    [Fact]
    public async Task Request_honors_message_envelopes()
    {
        await using var system = new ActorSystem();

        var echo = Props.FromFunc(ctx =>
            {
                if (ctx.Sender is not null)
                {
                    if (ctx.Headers.Count == 1)
                    {
                        ctx.Respond(ctx.Headers["foo"]);
                    }
                    else
                    {
                        ctx.Respond("Invalid headers in request");
                    }
                }

                return Task.CompletedTask;
            }
        );

        var pid = system.Root.Spawn(echo);

        var wrongPid = PID.FromAddress("some-incorrect-address", "some-id");

        using var future = system.Future.Get();

        system.Root.Request(pid,
            new MessageEnvelope(1, wrongPid, MessageHeader.Empty.With("foo", "bar")),
            future.Pid
        );

        var response = await future.Task;

        response.Should().Be("bar");
    }
    
    [Fact]
    public async Task Actor_Request_honors_message_envelopes()
    {
        await using var system = new ActorSystem();

        var echo = Props.FromFunc(ctx =>
            {
                if (ctx.Sender is not null)
                {
                    if (ctx.Headers.Count == 1)
                    {
                        ctx.Request(ctx.Sender, MessageEnvelope.Wrap(ctx.Message!, ctx.Headers));
                    }
                    else
                    {
                        ctx.Respond("Invalid headers in request");
                    }
                }

                return Task.CompletedTask;
            }
        );

        var pid = system.Root.Spawn(echo);

        var wrongPid = PID.FromAddress("some-incorrect-address", "some-id");

        using var future = system.Future.Get();

        system.Root.Request(pid,
            new MessageEnvelope(1, wrongPid, MessageHeader.Empty.With("foo", "bar")),
            future.Pid
        );

        var response = await future.Task;

        response.Should().BeOfType<MessageEnvelope>()
            .Which.Header["foo"].Should().Be("bar");
    }

    public record SomeRequest;

    public record SomeResponse;

    public record StartMessage;
}